from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.utils.security import get_current_user, verify_token
from app.schemas.user import UserOut
from app.services.ws_handler import WebSocketHandler
from fastapi.security.utils import get_authorization_scheme_param
import traceback
from loguru import logger
import asyncio

router = APIRouter(prefix="/ws", tags=["WebSocket"])  # 修改前缀为 /ws

# 创建WebSocket处理器实例
ws_handler = WebSocketHandler()

@router.websocket("/{token}")
async def websocket_endpoint(
    websocket: WebSocket, 
    token: str
):
    user = None
    connection_established = False
    
    try:
        # 确保心跳检测任务已启动
        await ws_handler.setup()
        
        # 验证token
        user = await verify_token(token)
        if not user:
            logger.error(f"WebSocket连接尝试，但token无效: {token[:10]}...")
            await websocket.close(code=4001, reason="Invalid token")
            return
            
        # 确保使用数字主键
        user_id = int(user.id)  # 确保转换为整数
        logger.debug(f"用户 {user_id} (业务ID: {user.userid}) WebSocket连接开始")
        
        # 连接WebSocket
        await ws_handler.connect(user_id, websocket)
        connection_established = True
        logger.success(f"用户 {user_id} WebSocket连接成功")
        
        # 确保WebSocket已经连接
        if not connection_established:
            logger.error(f"用户 {user_id} WebSocket连接失败，跳过后续处理")
            return
            
        # 在连接建立后，主动尝试发送一次离线消息
        try:
            from app.services.message import MessageService
            await MessageService.retry_pending_messages(ws_handler, user_id)
        except Exception as e:
            logger.error(f"发送离线消息失败: {str(e)}")
        
        try:
            # 发送一个初始的ping消息，确认连接正常
            try:
                await websocket.send_json({"type": 0, "body": "ping"})
            except Exception as e:
                logger.error(f"发送ping消息失败，连接可能已断开: {str(e)}")
                return
                
            while True:
                try:
                    # 接收消息
                    data = await websocket.receive_json()
                    msg_type = data.get("type", "未知")
                    logger.debug(f"接收到用户 {user_id} 的消息")
                    
                    # 特别处理群聊消息，确保正确转发
                    if msg_type == 2:  # 群聊消息
                        logger.info(f"收到群聊消息，用户: {user_id}, 房间: {data.get('roomId')}")
                    
                    await ws_handler.handle_message(user_id, data)
                except WebSocketDisconnect:
                    logger.info(f"用户 {user_id} WebSocket连接断开")
                    break
                except Exception as e:
                    logger.error(f"处理用户 {user_id} 消息时发生错误: {str(e)}")
                    if "connection closed" in str(e).lower() or "not connected" in str(e).lower():
                        logger.info(f"用户 {user_id} 连接已关闭或未连接，中断消息循环")
                        break
                    print(traceback.format_exc())
                    continue
                    
        except Exception as e:
            logger.error(f"用户 {user_id} WebSocket连接错误: {str(e)}")
            print(traceback.format_exc())
            
    except Exception as e:
        client_host = websocket.client.host if hasattr(websocket, 'client') else "未知"
        if user:
            logger.error(f"用户 {user.id} ({client_host}) WebSocket端点错误: {str(e)}")
        else:
            logger.error(f"未认证用户 ({client_host}) WebSocket端点错误: {str(e)}")
        print(traceback.format_exc())
        
    finally:
        # 安全断开连接
        if user and connection_established:
            try:
                await ws_handler.disconnect(int(user.id))  # 确保使用数字主键
                logger.success(f"用户 {user.id} 已从WebSocket断开连接")
            except Exception as e:
                logger.error(f"断开用户 {user.id} 连接时发生错误: {str(e)}")
                print(traceback.format_exc())